import axios from 'axios';
import FormData from 'form-data';
import fs from 'fs';
import { v4 as uuidv4 } from 'uuid';
import { updateWorkflowWithImage } from '@/utils/comfyUIWorkflow';

const COMFYUI_API_URL = process.env.COMFYUI_API_URL;

export async function uploadImageToComfyUI(filePath) {
  const formData = new FormData();
  formData.append('image', fs.createReadStream(filePath));
  formData.append('overwrite', 'true');
  formData.append('type', 'input');

  try {
    const response = await axios.post(`${COMFYUI_API_URL}/upload/image`, formData, {
      headers: {
        ...formData.getHeaders(),
      },
    });
    console.log("uploadImageToComfyUI response:", response.data);
    return response.data;
  } catch (error) {
    console.error("uploadImageToComfyUI error:", error);
    throw error;
  }
}

export async function queuePrompt(workflowName, imageUrls) {
  try {
    const clientId = uuidv4();
    let workflow = require(`@/workflow/${workflowName}`);
    const workflowConfig = require(`@/workflow/config.json`);
    const imageNodes = workflowConfig[workflowName].LoadImage;

    if (imageNodes.length !== imageUrls.length) {
      throw new Error("Image count does not match workflow requirements");
    }

    workflow = updateWorkflowWithImage({
      workflow,
      imageNodes,
      imagePathArr: imageUrls,
    });

    const response = await axios.post(`${COMFYUI_API_URL}/prompt`, {
      client_id: clientId,
      prompt: workflow,
    });

    console.log("queuePrompt response:", response.data);
    const { prompt_id: promptId, error, number, node_errors: nodeErrors } = response.data;
    return { promptId, error, clientId, number, nodeErrors };
  } catch (error) {
    console.error("queuePrompt error:", error);
    throw error;
  }
}

export async function getPromptHistory(promptId) {
  try {
    const queueResponse = await axios.get(`${COMFYUI_API_URL}/queue`);
    if (taskIsPending(queueResponse.data?.queue_pending, promptId)) {
      return { status: "pending" };
    }

    const historyResponse = await axios.get(`${COMFYUI_API_URL}/history/${promptId}`);
    if (!historyResponse.data[promptId]) {
      return { status: "pending" };
    }

    const { completed, status_str } = historyResponse.data[promptId].status;
    if (completed && status_str === "success") {
      const images = extractAllImagesFromResponse(promptId, historyResponse);
      return {
        status: "success",
        images: images.map(image => `${COMFYUI_API_URL}/view?filename=${image.filename}&type=${image.type}&subfolder=${image.subfolder}`),
        message: "success",
      };
    } else if (status_str === "error") {
      return {
        status: "error",
        message: historyResponse.data[promptId].status.messages,
      };
    } else {
      return { status: "pending" };
    }
  } catch (error) {
    console.error("getPromptHistory error:", error);
    throw error;
  }
}

function taskIsPending(pendingTasks, promptId) {
  if (!pendingTasks) return false;
  return pendingTasks.some(task => task[1] === promptId);
}

function extractAllImagesFromResponse(promptId, response) {
  const workflow = response.data[promptId].prompt[2];
  const saveImageNodeIds = parseWorkflow(workflow, "SaveImage");
  const outputs = response.data[promptId].outputs;
  return saveImageNodeIds
    .flatMap(node => outputs[node].images)
    .filter(image => image.type === "output");
}

// 这个函数需要在 utils/comfyUIWorkflow.js 中实现
function parseWorkflow(workflow, nodeType) {
  // 实现解析 workflow 并返回指定类型节点的 ID 的逻辑
}